import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import time
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline

# ==================== 初始化设置 ====================
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

# 回测参数
initial_capital = 1_000_000
start_date = '20180101'
end_date = '20201231'
stock_count = 10  # 测试用股票数量
lookback_window = 60  # 特征回溯窗口
forecast_horizon = 5  # 预测未来5日收益


# ==================== 数据准备 ====================
def compute_rsi(series, period=14):
    """计算RSI指标"""
    delta = series.diff().dropna()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)

    avg_gain = gain.rolling(period).mean()
    avg_loss = loss.rolling(period).mean()

    rs = avg_gain / avg_loss
    return 100 - (100 / (1 + rs))


def prepare_ml_data_with_retry():
    """带重试机制的数据准备函数"""
    print("正在准备数据...")
    max_retry = 3
    all_data = {}

    # 获取沪深300成分股
    for _ in range(max_retry):
        try:
            df_component = pro.index_weight(index_code='399300.SZ',
                                            start_date=start_date,
                                            end_date=end_date)
            stock_list = df_component['con_code'].unique().tolist()[:stock_count]
            break
        except Exception as e:
            print(f"获取成分股失败，重试中... 错误：{str(e)}")
            time.sleep(2)

    # 获取行情数据并生成特征
    for ts_code in tqdm(stock_list, desc="下载行情数据"):
        for retry in range(max_retry):
            try:
                df = pro.daily(ts_code=ts_code,
                               start_date=start_date,
                               end_date=end_date)
                if not df.empty:
                    # 数据预处理
                    df = df.sort_values('trade_date').reset_index(drop=True)
                    df['trade_date'] = df['trade_date'].astype(str)

                    # 计算技术指标
                    df['returns'] = df['close'].pct_change()
                    df['MA5'] = df['close'].rolling(5).mean()
                    df['MA20'] = df['close'].rolling(20).mean()
                    df['Volatility'] = df['returns'].rolling(20).std()
                    df['RSI'] = compute_rsi(df['close'])

                    # 添加市场数据
                    index_data = pro.index_daily(ts_code='399300.SZ',
                                                 start_date=start_date,
                                                 end_date=end_date)[['trade_date', 'close']]
                    index_data = index_data.rename(columns={'close': 'Index_Close'})
                    index_data['trade_date'] = index_data['trade_date'].astype(str)
                    df = pd.merge(df, index_data, on='trade_date', how='left')

                    # 生成双均线信号
                    df['signal'] = 0
                    df['prev_MA5'] = df['MA5'].shift(1)
                    df['prev_MA20'] = df['MA20'].shift(1)
                    golden_cross = (df['prev_MA5'] <= df['prev_MA20']) & (df['MA5'] > df['MA20'])
                    death_cross = (df['prev_MA5'] >= df['prev_MA20']) & (df['MA5'] < df['MA20'])
                    df.loc[golden_cross, 'signal'] = 1
                    df.loc[death_cross, 'signal'] = -1
                    df['position'] = df['signal'].replace(0, np.nan).ffill().fillna(0)
                    df['return'] = df['close'].pct_change()

                    # 创建标签
                    df['label'] = (df['close'].shift(-forecast_horizon) / df['close'] - 1)
                    df['label'] = np.where(df['label'] > 0, 1, 0)

                    df = df.dropna()
                    all_data[ts_code] = df
                break
            except Exception as e:
                print(f"获取{ts_code}数据失败，重试{retry + 1}/{max_retry}")
                time.sleep(1)

    return all_data, stock_list


# ==================== 机器学习模型 ====================
class MLStrategy:
    def __init__(self):
        self.model = Pipeline([
            ('scaler', StandardScaler()),
            ('rf', RandomForestClassifier(n_estimators=100,
                                          max_depth=5,
                                          random_state=42))
        ])

    def train_predict(self, train_data, test_features):
        X_train = train_data.drop(['label', 'trade_date'], axis=1)
        y_train = train_data['label']
        self.model.fit(X_train, y_train)
        X_test = test_features.drop(['trade_date'], axis=1)
        return self.model.predict_proba(X_test)[:, 1]


# ==================== 回测引擎 ====================
def enhanced_backtest(all_data, stock_list):
    """双均线+机器学习混合回测"""
    dates = sorted(set([d for df in all_data.values() for d in df['trade_date']]))
    print(f"\n总交易日数: {len(dates)}")

    # 初始化策略
    cash_ma, cash_ml = initial_capital, initial_capital
    holdings_ma, holdings_ml = {}, {}
    portfolio_values_ma, portfolio_values_ml = [], []

    # 生成调仓日期
    rebalance_dates = pd.date_range(start=pd.to_datetime(start_date),
                                    end=pd.to_datetime(end_date),
                                    freq='QE') + pd.offsets.MonthBegin(1)
    rebalance_dates = [d.strftime('%Y%m%d') for d in rebalance_dates]
    valid_rebalance_dates = [d for d in rebalance_dates if d in dates]

    # 机器学习参数
    ml_engine = MLStrategy()
    feature_cols = ['open', 'high', 'low', 'close', 'vol',
                    'MA5', 'MA20', 'Volatility', 'RSI', 'Index_Close']

    for i, date in enumerate(tqdm(dates, desc="回测进度")):
        current_date_str = date

        # ===== 双均线策略 =====
        # 止损止盈逻辑
        to_delete = []
        for code in list(holdings_ma.keys()):
            try:
                df = all_data[code]
                day_data = df[df['trade_date'] == current_date_str]
                if not day_data.empty:
                    close_price = day_data['close'].values[0]
                    entry_price = holdings_ma[code]['entry_price']
                    highest_price = holdings_ma[code]['highest_price']

                    current_highest = max(highest_price, close_price)
                    drawdown = (current_highest - close_price) / current_highest
                    profit = (close_price - entry_price) / entry_price

                    if drawdown >= 0.1 or profit >= 0.2:
                        cash_ma += holdings_ma[code]['shares'] * close_price
                        to_delete.append(code)
                    else:
                        holdings_ma[code]['highest_price'] = current_highest
            except Exception as e:
                print(f"MA平仓异常 {code}: {str(e)}")

        for code in to_delete:
            del holdings_ma[code]

        # 买入信号处理
        for code in stock_list:
            try:
                df = all_data[code]
                day_data = df[df['trade_date'] == current_date_str]
                if not day_data.empty and day_data['signal'].iloc[0] == 1:
                    if code not in holdings_ma and cash_ma > 10000:
                        close_price = day_data['close'].iloc[0]
                        shares = int((cash_ma * 0.1) // close_price)
                        if shares > 0:
                            holdings_ma[code] = {
                                'shares': shares,
                                'entry_price': close_price,
                                'highest_price': close_price
                            }
                            cash_ma -= shares * close_price
            except Exception as e:
                print(f"MA买入异常 {code}: {str(e)}")

        # ===== 机器学习策略 =====
        if current_date_str in valid_rebalance_dates:
            # 清仓
            for code in list(holdings_ml.keys()):
                try:
                    df = all_data[code]
                    day_data = df[df['trade_date'] == current_date_str]
                    if not day_data.empty:
                        cash_ml += holdings_ml[code] * day_data['close'].iloc[0]
                except:
                    pass
            holdings_ml.clear()

            # 准备训练数据
            train_dates = dates[max(0, i - lookback_window):i]
            all_train = []
            for code in stock_list:
                df = all_data.get(code, pd.DataFrame())
                if df is not None and not df.empty:
                    train_data = df[df['trade_date'].isin(train_dates)]
                    if len(train_data) > 20:
                        all_train.append(train_data[feature_cols + ['label', 'trade_date']])
            if not all_train:
                continue
            full_train = pd.concat(all_train)

            # 准备测试数据
            test_features = []
            valid_codes = []
            for code in stock_list:
                df = all_data.get(code, pd.DataFrame())
                if df is not None and not df.empty:
                    idx = df.index[df['trade_date'] == current_date_str]
                    if len(idx) > 0 and idx[0] + 1 < len(df):
                        test_data = df.iloc[idx[0] + 1]
                        test_features.append(test_data[feature_cols + ['trade_date']])
                        valid_codes.append(code)
            if not test_features:
                continue
            test_df = pd.DataFrame(test_features)

            # 生成预测
            try:
                probabilities = ml_engine.train_predict(full_train, test_df)
                predictions = pd.Series(probabilities, index=valid_codes)
                predictions = predictions.sort_values(ascending=False).head(3)
            except:
                predictions = pd.Series()

            # 执行买入
            if not predictions.empty:
                per_cash = cash_ml / len(predictions)
                for code, prob in predictions.items():
                    try:
                        df = all_data[code]
                        day_data = df[df['trade_date'] == current_date_str]
                        if not day_data.empty:
                            price = day_data['close'].iloc[0]
                            shares = int(per_cash // price)
                            if shares > 0:
                                holdings_ml[code] = shares
                                cash_ml -= shares * price
                    except Exception as e:
                        print(f"ML买入异常 {code}: {str(e)}")

        # 计算净值
        total_ma = cash_ma + sum([
            info['shares'] * all_data[code][all_data[code]['trade_date'] == current_date_str]['close'].iloc[0]
            for code, info in holdings_ma.items()
            if not all_data[code][all_data[code]['trade_date'] == current_date_str].empty
        ])

        total_ml = cash_ml + sum([
            shares * all_data[code][all_data[code]['trade_date'] == current_date_str]['close'].iloc[0]
            for code, shares in holdings_ml.items()
            if not all_data[code][all_data[code]['trade_date'] == current_date_str].empty
        ])

        portfolio_values_ma.append(total_ma)
        portfolio_values_ml.append(total_ml)

    return portfolio_values_ma, portfolio_values_ml, dates


# ==================== 绩效分析 ====================
def enhanced_analysis(values, dates, name):
    series = pd.Series(values, index=pd.to_datetime(dates))
    returns = series.pct_change().fillna(0)

    cum_return = (1 + returns).cumprod()
    total_return = cum_return.iloc[-1] - 1
    annual_return = (1 + total_return) ** (252 / len(returns)) - 1
    volatility = returns.std() * np.sqrt(252)
    sharpe = annual_return / volatility if volatility != 0 else np.nan
    max_dd = (cum_return.cummax() - cum_return).max()

    print(f"\n{name}策略绩效:")
    print(f"累计收益率: {total_return:.2%}")
    print(f"年化收益率: {annual_return:.2%}")
    print(f"波动率: {volatility:.2%}")
    print(f"夏普比率: {sharpe:.2f}")
    print(f"最大回撤: {max_dd:.2%}")

    return cum_return


# ==================== 主程序 ====================
if __name__ == "__main__":
    # 数据准备
    all_data, stock_list = prepare_ml_data_with_retry()
    valid_stocks = [k for k in all_data.keys() if not all_data[k].empty]
    print(f"\n有效股票数量: {len(valid_stocks)}/{len(stock_list)}")

    # 执行回测
    values_ma, values_ml, dates = enhanced_backtest(all_data, valid_stocks)

    # 绩效分析
    cum_ma = enhanced_analysis(values_ma, dates, "双均线")
    cum_ml = enhanced_analysis(values_ml, dates, "机器学习")

    # 可视化
    plt.figure(figsize=(12, 6))
    cum_ma.plot(label='Dual MA Strategy')
    cum_ml.plot(label='Machine Learning Strategy')
    plt.title('Strategy Performance Comparison')
    plt.xlabel('Date')
    plt.ylabel('Portfolio Value')
    plt.legend()
    plt.grid(True)
    plt.show()
